#!/usr/bin/env python3
"""
MySQL MCP Server 主程序入口

本模块是 MySQL MCP Server 的主程序入口，负责：
- 解析命令行参数
- 初始化服务器配置
- 启动 MCP 服务器
- 处理服务器生命周期
- 信号处理和优雅关闭
"""

import asyncio
import signal
import sys
import os
import argparse
from typing import Optional
from pathlib import Path

from mysql_mcp_server.config import DatabaseConfig, load_config_from_env
from mysql_mcp_server.mcp_server import MCPServer, create_mcp_server
from mysql_mcp_server.logger import get_logger, setup_logger, mcp_logger
from mysql_mcp_server.query_engine import create_query_engine


def setup_signal_handlers(shutdown_event: asyncio.Event) -> None:
    """
    设置信号处理器
    
    处理 SIGINT 和 SIGTERM 信号以实现优雅关闭。
    
    Args:
        shutdown_event: 关闭事件
    """
    def signal_handler(signum, frame):
        logger = get_logger(__name__)
        logger.info(f"接收到信号 {signum}，开始优雅关闭...")
        shutdown_event.set()
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)


class MCPServerApplication:
    """
    MCP 服务器应用程序
    
    管理服务器的完整生命周期，包括启动、运行和关闭。
    """
    
    def __init__(self):
        """
        初始化应用程序
        """
        self.logger = get_logger(__name__)
        self.server: Optional[MCPServer] = None
        self.shutdown_event = asyncio.Event()
        self.running = False
    
    def setup_signal_handlers(self) -> None:
        """
        设置信号处理器
        
        处理 SIGINT 和 SIGTERM 信号以实现优雅关闭。
        """
        def signal_handler(signum, frame):
            self.logger.info(f"接收到信号 {signum}，开始优雅关闭...")
            self.shutdown_event.set()
        
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)
    
    async def start_server(self, db_config: DatabaseConfig, 
                          enable_write_operations: bool = False) -> None:
        """
        启动 MCP 服务器
        
        Args:
            db_config: 数据库配置
            enable_write_operations: 是否启用写操作
        """
        try:
            # 创建并启动服务器
            self.server = await create_mcp_server(db_config)
            
            # 创建查询引擎并集成到服务器
            query_engine = create_query_engine(
                db_pool=self.server.db_pool,
                enable_write_operations=enable_write_operations,
                max_result_size=int(os.getenv('MCP_MAX_RESULT_SIZE', '10000')),
                query_timeout=int(os.getenv('MCP_QUERY_TIMEOUT', '30'))
            )
            
            # 将查询引擎集成到服务器
            self.server.query_engine = query_engine
            
            self.running = True
            self.logger.info("MCP 服务器启动成功，等待客户端连接...")
            
        except Exception as e:
            self.logger.error("启动服务器失败", error=str(e))
            raise
    
    async def run_stdio_server(self) -> None:
        """
        运行基于 stdio 的 MCP 服务器
        
        通过标准输入输出与客户端通信。
        """
        if not self.server:
            raise RuntimeError("服务器未初始化")
        
        self.logger.info("启动 stdio MCP 服务器")
        
        try:
            # 创建异步任务处理输入输出
            input_task = asyncio.create_task(self._handle_stdin())
            shutdown_task = asyncio.create_task(self._wait_for_shutdown())
            
            # 等待任一任务完成
            done, pending = await asyncio.wait(
                [input_task, shutdown_task],
                return_when=asyncio.FIRST_COMPLETED
            )
            
            # 取消未完成的任务
            for task in pending:
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
            
        except Exception as e:
            self.logger.error("运行服务器时发生错误", error=str(e))
            raise
    
    async def _handle_stdin(self) -> None:
        """
        处理标准输入
        
        从标准输入读取 JSON-RPC 请求并处理。
        """
        reader = asyncio.StreamReader()
        protocol = asyncio.StreamReaderProtocol(reader)
        await asyncio.get_event_loop().connect_read_pipe(
            lambda: protocol, sys.stdin
        )
        
        while self.running and not self.shutdown_event.is_set():
            try:
                # 读取一行输入
                line = await reader.readline()
                if not line:
                    break
                
                request_data = line.decode('utf-8').strip()
                if not request_data:
                    continue
                
                # 处理请求
                response_data = await self.server.handle_request(request_data)
                
                # 输出响应
                print(response_data, flush=True)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error("处理输入时发生错误", error=str(e))
                # 继续处理下一个请求
                continue
    
    async def _wait_for_shutdown(self) -> None:
        """
        等待关闭信号
        """
        await self.shutdown_event.wait()
        self.running = False
    
    async def stop_server(self) -> None:
        """
        停止服务器
        
        执行优雅关闭流程。
        """
        if self.server:
            try:
                await self.server.stop()
                self.logger.info("服务器已成功停止")
            except Exception as e:
                self.logger.error("停止服务器时发生错误", error=str(e))
        
        self.running = False


def parse_arguments() -> argparse.Namespace:
    """
    解析命令行参数
    
    Returns:
        argparse.Namespace: 解析后的参数
    """
    parser = argparse.ArgumentParser(
        description="MySQL MCP Server - 为 Model Context Protocol 提供 MySQL 数据库访问",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  %(prog)s --host localhost --port 3306 --database mydb --user myuser
  %(prog)s --config-file /path/to/config.env
  %(prog)s --enable-write-operations --log-level DEBUG

环境变量:
  MYSQL_HOST, MYSQL_PORT, MYSQL_DATABASE, MYSQL_USER, MYSQL_PASSWORD
  MCP_LOG_LEVEL, MCP_LOG_FORMAT, MCP_ENABLE_WRITE_OPERATIONS
        """
    )
    
    # 数据库连接参数
    db_group = parser.add_argument_group('数据库连接')
    db_group.add_argument(
        '--host', '-H',
        help='MySQL 服务器主机名 (默认: localhost)'
    )
    db_group.add_argument(
        '--port', '-P',
        type=int,
        help='MySQL 服务器端口 (默认: 3306)'
    )
    db_group.add_argument(
        '--database', '-d',
        help='数据库名称'
    )
    db_group.add_argument(
        '--user', '-u',
        help='数据库用户名'
    )
    db_group.add_argument(
        '--password', '-p',
        help='数据库密码（不推荐，请使用环境变量）'
    )
    
    # 服务器配置
    server_group = parser.add_argument_group('服务器配置')
    server_group.add_argument(
        '--enable-write-operations',
        action='store_true',
        help='启用写操作 (INSERT, UPDATE, DELETE 等)'
    )
    server_group.add_argument(
        '--config-file',
        type=Path,
        help='配置文件路径 (.env 格式)'
    )
    
    # 日志配置
    log_group = parser.add_argument_group('日志配置')
    log_group.add_argument(
        '--log-level',
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
        help='日志级别 (默认: INFO)'
    )
    log_group.add_argument(
        '--log-format',
        choices=['console', 'json'],
        help='日志格式 (默认: console)'
    )
    
    # 其他选项
    parser.add_argument(
        '--version',
        action='version',
        version='MySQL MCP Server 1.0.0'
    )
    parser.add_argument(
        '--test-connection',
        action='store_true',
        help='测试数据库连接后退出'
    )
    
    return parser.parse_args()


def load_configuration(args: argparse.Namespace) -> DatabaseConfig:
    """
    加载配置
    
    Args:
        args: 命令行参数
        
    Returns:
        DatabaseConfig: 数据库配置
    """
    # 如果指定了配置文件，先加载它
    if args.config_file:
        if not args.config_file.exists():
            raise FileNotFoundError(f"配置文件不存在: {args.config_file}")
        
        # 加载 .env 文件
        with open(args.config_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip()
    
    # 从环境变量加载基础配置
    config = load_config_from_env()
    
    # 命令行参数覆盖环境变量
    if args.host:
        config.host = args.host
    if args.port:
        config.port = args.port
    if args.database:
        config.database = args.database
    if args.user:
        config.user = args.user
    if args.password:
        config.password = args.password
    
    return config


async def test_database_connection(config: DatabaseConfig) -> bool:
    """
    测试数据库连接
    
    Args:
        config: 数据库配置
        
    Returns:
        bool: 连接是否成功
    """
    logger = get_logger(__name__)
    
    try:
        from .connection_pool import DatabaseConnectionPool
        
        logger.info("测试数据库连接...")
        logger.info(f"连接到: {config.host}:{config.port}/{config.database}")
        
        # 创建连接池
        pool = DatabaseConnectionPool(config)
        pool.initialize()
        
        # 执行测试查询
        result = pool.execute_query("SELECT 1 as test")
        
        if result and len(result) > 0 and result[0].get('test') == 1:
            logger.info("数据库连接测试成功！")
            pool.close()
            return True
        else:
            logger.error("数据库连接测试失败：查询结果异常")
            pool.close()
            return False
            
    except Exception as e:
        logger.error(f"数据库连接测试失败: {str(e)}")
        return False


async def main() -> int:
    """
    主函数
    
    Returns:
        int: 退出代码
    """
    try:
        # 解析命令行参数
        args = parse_arguments()
        
        # 设置日志
        log_level = args.log_level or os.getenv('MCP_LOG_LEVEL', 'INFO')
        log_format = args.log_format or os.getenv('MCP_LOG_FORMAT', 'console')
        setup_logger(log_level, log_format == 'json')
        
        logger = get_logger(__name__)
        logger.info("启动 MySQL MCP Server")
        
        # 加载配置
        try:
            config = load_configuration(args)
        except Exception as e:
            logger.error(f"加载配置失败: {str(e)}")
            return 1
        
        # 验证配置
        validation_errors = config.validate()
        if validation_errors:
            logger.error("配置验证失败:")
            for error in validation_errors:
                logger.error(f"  - {error}")
            return 1
        
        # 如果是测试连接模式
        if args.test_connection:
            success = await test_database_connection(config)
            return 0 if success else 1
        
        # 创建应用程序
        app = MCPServerApplication()
        app.setup_signal_handlers()
        
        # 启动服务器
        enable_write_ops = (
            args.enable_write_operations or 
            os.getenv('MCP_ENABLE_WRITE_OPERATIONS', 'false').lower() == 'true'
        )
        
        await app.start_server(config, enable_write_ops)
        
        # 运行服务器
        await app.run_stdio_server()
        
        # 停止服务器
        await app.stop_server()
        
        logger.info("MySQL MCP Server 已退出")
        return 0
        
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在退出...")
        return 0
    except Exception as e:
        logger.error(f"程序运行时发生未处理的错误: {str(e)}")
        return 1


def run_server() -> None:
    """
    运行服务器的入口函数
    
    这个函数可以被 setup.py 的 entry_points 调用。
    """
    try:
        exit_code = asyncio.run(main())
        sys.exit(exit_code)
    except KeyboardInterrupt:
        sys.exit(0)
    except Exception as e:
        print(f"Fatal error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    run_server()